import numpy as np
import XmlTools
from . import FSTools
from multiprocessing import Pool

try:
    import pyfits as pf
except ImportError:
    from astropy.io import fits as pf


"""
try:
    from scipy import interpolate, optimize
    from scipy.integrate import quad
except ImportError:
    def collect_containment_radius(psffile, **kwds):
        print '[ImportError] Cannot import scipy!'
        if isinstance(psffile, list) or isinstance(psffile, tuple):
            fh = open('psffile_temp.lst', 'w')
            for psff in psffile:
                fh.write(psff+'\n')
            fh.close()
            print '[collect_containment_radius] Save psffile to psffile_temp.lst!'
else:
    def simple_containment_radius(psffile, cdf=0.95, outfile=None):
        hdulist = pf.open(psffile, memmap=True)
        theta = hdulist[2].data.field(0) # [deg]
        energy = hdulist[1].data.field(0) # [MeV]
        psf = hdulist[1].data.field(2)
        hdulist.close()
        del hdulist[1].data, hdulist[2].data, hdulist

        _cdf_precalc = cdf/(2.*np.pi)
        d2r = np.pi/180.
        x_val = theta * d2r # rad
        ys_val = x_val.reshape(1, -1) * psf # theta*psf(theta, E)
        max_x_val = x_val.max()

        radius = []
        for iE in range(ys_val.shape[0]):
            intpf = interpolate.interp1d(x_val, ys_val[iE, ...], bounds_error=True)
            func_to_solve = lambda thmax: quad(intpf, 0, thmax)[0] - _cdf_precalc
            _r = optimize.brentq(func_to_solve, 0., max_x_val, disp=False, full_output=False)
            radius.append(_r)

        radius = np.atleast_1d(radius)/d2r
        if outfile is None:
            return energy, radius
        else:
            outdata = np.c_[energy, radius]
            np.savetxt(outfile, outdata, delimiter='\t',
                    header='E[MeV]\tRadius[deg]\t##contain=%f'%cdf)

    def collect_containment_radius(psffile_list, cdf=0.95, outfile=None):
        if isinstance(psffile_list, str):
            fh = open(psffile_list, 'r')
            psffiles = [l.strip() for l in fh]
            fh.close()
        elif isinstance(psffile_list, list) or isinstance(psffile_list, tuple):
            psffiles = psffile_list

        outdata = []
        for psffile in psffiles:
            _e, _rad = simple_containment_radius(psffile, cdf=cdf, outfile=None)
            outdata.append(_rad)
        outdata = np.atleast_2d(outdata)

        if outfile is None:
            return _e, psffiles, outdata
        else:
            output_data = np.c_[psffiles, outdata]
            elines = _e.tolist()
            elines.insert(0, 'ENERGY')
            output_data = np.r_[np.atleast_2d(elines), output_data]
            np.savetxt(outfile, output_data, delimiter='\t', fmt='%s',
                    header='ITEM\tRadius[deg]\t##contain=%f'%cdf)
"""

def get_psf_according_to_model(xmlfile, ltcube, irfs, evtype, emin, emax, nenergies,
                               masksrcs=None, thetamax=90., ntheta=900, outdir=None,
                               njobs=1, psflist=None):
    x = XmlTools.SMKernel.SourceModel(xmlfile, raiseErr=True)
    if isinstance(masksrcs, str):
        is_srcmask = True
        fh = open(masksrcs, 'r')
        masksrcs = [s.strip() for s in fh]
        fh.close()
    elif isinstance(masksrcs, list) or isinstance(masksrcs, tuple):
        is_srcmask = True
    else:
        is_srcmask = False

    if outdir is None or (not outdir):
        outfile_prefix = './'
    elif outdir.endswith('/'):
        outfile_prefix = outdir
    else:
        outfile_prefix = outdir + '/'

    mykwds, psffiles, ras, decs = [], [], [], []
    for sname, src in x.srcList.items():
        if is_srcmask and (not sname in masksrcs):
            continue
        ra = src.spatialModel.RA.value
        dec = src.spatialModel.DEC.value

        outfile = outfile_prefix + 'mypsf_%s.fits'%'_'.join(sname.split())
        kwds = {'ltcube' : ltcube, 'irfs' : irfs, 'evtype' : evtype, 'ra' : ra,
                'dec' : dec, 'emin' : emin, 'emax' : emax, 'nenergies' : nenergies,
                'outfile' : outfile, 'thetamax' : thetamax, 'ntheta' : ntheta}
        mykwds.append(kwds)
        psffiles.append(outfile)
        ras.append(ra)
        decs.append(dec)

    if njobs < 2:
        for kwds in mykwds:
            FSTools.gtpsf(**kwds)
    else:
        pool = Pool(processes=njobs)
        for kwds in mykwds:
            pool.apply_async(FSTools._gtpsf_mp_drv, (kwds, ))
        pool.close()
        pool.join()
    if not psflist is None:
        outdata = np.c_[psffiles, ras, decs]
        np.savetxt(psflist, outdata, delimiter='\t', fmt='%s', header='FileName\tRA[deg]\tDecl[deg]')
    return psffiles, ras, decs

def parse_covardat(infile, outfile=None):
    """
    covert the covariance dat file into a dictionary or npz file
    """
    covs = []
    pnamesmap = dict()
    state = 0 # flag of FSM
    with open(infile, 'r') as fh:
        for l in fh:
            if l.startswith('---'):
                state = 0

            if state == 1:
                # deal with pnames
                lspl = l.strip().split()
                pidx = int(lspl[0])
                pname = ' '.join(lspl[1:])
                pnamesmap[pidx] = pname
            elif state == 2:
                # deal with cov mat
                lspl = l.strip().split()
                covs.append([float(v) for v in lspl])
            else:
                if '--free parameters--' in l:
                    state = 1
                elif '--covariance matrix--' in l:
                    state = 2

    covmat = np.array(covs)
    pnames = [pnamesmap[i] for i in range(1,len(pnamesmap)+1)]
    sig = np.sqrt(np.diagonal(covmat))
    cormat = covmat/(sig.reshape(-1,1)*sig.reshape(1,-1))
    outdict = dict(pnames=pnames, corr=cormat, cov=covmat)

    if outfile is None:
        return outdict
    else:
        np.savez_compressed(outfile, **outdict)
